import numpy as np
import pandas as pd
import os
import glob

# 设定数据目录和输出目录
input_directory = r'D:\pythonProject\分类识别\测试集'  # 输入目录，存放原始npy文件
output_directory = r'D:\pythonProject\分类识别\测试集_副本'  # 输出目录，存放拆分后的npy文件

# 确保输出目录存在
os.makedirs(output_directory, exist_ok=True)

# 遍历输入目录中的.npy文件
for file_path in glob.glob(os.path.join(input_directory, '*.npy')):
    # 加载数据
    data = np.load(file_path, allow_pickle=True)

    # 转换为 DataFrame 以便于处理
    df = pd.DataFrame(data)

    # 假设 segment_id 是数据中的列名称，修改为实际的列索引或名称
    if 'segment_id' not in df.columns:
        print(f"Warning: 'segment_id' not found in {file_path}. Skipping.")
        continue

    # 识别独特的航迹
    unique_segments = df['segment_id'].unique()

    # 遍历每个唯一的航迹，保存为单独的.npy文件
    for segment in unique_segments:
        # 筛选出当前航迹的数据
        segment_data = df[df['segment_id'] == segment]

        # 获取类名称，从segment_id中进行提取，如果需要前缀
        # 这里假设segment_id形式可以从中提取航迹类别
        # 如果需要直接使用 segment_id，您可以修改这一部分
        target_name = segment_data['icaoaircrafttype'].iloc[0]  # 假设有一个'target_name'列
        count = 1

        # 创建唯一文件名
        while True:
            filename = f"{target_name}({count}).npy"
            output_path = os.path.join(output_directory, filename)
            if not os.path.exists(output_path):
                break  # 找到一个唯一的文件名
            count += 1

        # 保存当前航迹
        np.save(output_path, segment_data.values)
        print(f"Saved {output_path}")

print("所有航迹已拆分保存。")